# -*- coding: utf-8 -*-
"""
主程序 - 基于YOLOv5的物体追踪系统
"""
import cv2
import time
from src.usb_camera import USBCamera
from src.yolo_detector import YOLODetector
from src.pan_tilt_controller import PanTiltController
from src.config import TRACKING_CONFIG, DISPLAY_CONFIG, PAN_TILT_CONFIG


class ObjectTracker:
    def __init__(self):
        """初始化追踪系统"""
        self.camera = USBCamera()
        # 使用YOLOv5模型（可替换为自定义训练模型）
        self.detector = YOLODetector()
        # 初始化云台控制器
        self.pan_tilt = PanTiltController()

        # 从config.py读取配置
        self.dead_zone = TRACKING_CONFIG['dead_zone']
        self.kp = TRACKING_CONFIG['kp']
        self.show_preview = DISPLAY_CONFIG['show_preview']
        self.print_detections = DISPLAY_CONFIG['print_detections']

        # 计算画面中心点 (基于写死的分辨率)
        self.frame_center_x = 640 // 2
        self.frame_center_y = 480 // 2

        print("物体追踪系统初始化完成!")

    def calculate_error_and_control(self, bbox):
        """
        计算目标与画面中心的偏差，并控制云台
        参数:
            bbox: 边界框 [x1, y1, x2, y2]
        """
        x1, y1, x2, y2 = bbox
        target_center_x = (x1 + x2) // 2
        target_center_y = (y1 + y2) // 2

        error_x = target_center_x - self.frame_center_x
        error_y = target_center_y - self.frame_center_y

        # 死区判断 - 在中心区域不移动
        if abs(error_x) < self.dead_zone and abs(error_y) < self.dead_zone:
            return error_x, error_y, (target_center_x, target_center_y)

        # 比例控制计算角度增量
        # 注意：这里的符号取决于摄像头安装方向和舵机转向
        pan_delta = -self.kp * error_x  # 水平方向，误差为正需要向左转
        tilt_delta = self.kp * error_y  # 垂直方向，误差为正需要向下转

        # 控制云台移动
        new_pan_angle, new_tilt_angle = self.pan_tilt.move_relative(pan_delta, tilt_delta)

        if self.print_detections:
            print(f"云台控制 - 水平: {new_pan_angle:.1f}°, 垂直: {new_tilt_angle:.1f}°")

        return error_x, error_y, (target_center_x, target_center_y)

    def run(self):
        """运行主追踪循环"""
        frame_count = 0
        start_time = time.time()

        print("开始追踪，按 'q' 退出...")

        try:
            while True:
                # 读取摄像头帧
                ret, frame = self.camera.read_frame()
                if not ret:
                    print("无法从摄像头读取帧")
                    break

                # YOLOv5目标检测
                detections, frame_with_bbox = self.detector.detect(frame)

                # 处理检测结果
                if detections:
                    # 选择置信度最高的检测结果
                    best_detection = max(detections, key=lambda x: x['confidence'])

                    # 计算偏差并控制云台
                    error_x, error_y, target_center = self.calculate_error_and_control(
                        best_detection['bbox']
                    )

                    # 绘制中心点和连线
                    cv2.circle(frame_with_bbox, (self.frame_center_x, self.frame_center_y),
                               5, (255, 0, 0), -1)  # 蓝色：画面中心
                    cv2.circle(frame_with_bbox, target_center, 5, (0, 0, 255), -1)  # 红色：目标中心
                    cv2.line(frame_with_bbox, (self.frame_center_x, self.frame_center_y),
                             target_center, (255, 255, 0), 2)  # 青色连线

                    # 显示偏差信息
                    cv2.putText(frame_with_bbox,
                                f"Error X: {error_x}, Y: {error_y}",
                                (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7,
                                (255, 255, 255), 2)

                    # 打印检测信息
                    if self.print_detections:
                        class_name = best_detection['class_name']
                        confidence = best_detection['confidence']
                        print(f"检测到: {class_name} 置信度: {confidence:.2f} 偏差: ({error_x}, {error_y})")
                else:
                    # 没有检测到目标
                    if self.print_detections and frame_count % 30 == 0:
                        print("未检测到目标")

                # 计算并显示帧率
                frame_count += 1
                if frame_count % 30 == 0:
                    fps = frame_count / (time.time() - start_time)
                    cv2.putText(frame_with_bbox, f"FPS: {fps:.1f}",
                                (10, 70), cv2.FONT_HERSHEY_SIMPLEX, 0.7,
                                (255, 255, 255), 2)
                    frame_count = 0
                    start_time = time.time()

                # 显示预览窗口
                if self.show_preview:
                    cv2.imshow('YOLOv5 Object Tracker', frame_with_bbox)

                # 检查退出按键
                key = cv2.waitKey(1) & 0xFF
                if key == ord('q'):
                    break

        except KeyboardInterrupt:
            print("\n程序被用户中断")
        finally:
            self.stop()

    def stop(self):
        """停止追踪系统"""
        self.pan_tilt.stop()  # 停止云台
        self.pan_tilt.close()  # 关闭串口
        self.camera.release()
        cv2.destroyAllWindows()
        print("追踪系统已停止")


if __name__ == "__main__":
    tracker = ObjectTracker()
    tracker.run()